# coding:utf-8
# Author : hiicy redldw
# Date : 2019/04/26
# FIXME: 待定
import tensorflow as tf
from keras.applications.mobilenetv2 import MobileNetV2
import functools
import pandas as pd


# 编写模型函数
def WeightVariable(shape):
    initial = tf.truncated_normal(shape, stddev=0.1)
    return tf.Variable(initial)


def BiasVariable(shape):
    initial = tf.constant(0.1, shape=shape)
    return tf.Variable(initial)


def Conv2d(input, W, b, conv_strides, padding, activation="relu"):
    x = tf.nn.conv2d(input, W, strides=conv_strides, padding=padding)
    x = tf.nn.relu6(x + b)
    x = tf.nn.batch_normalization(x, 0, variance=1)
    return x


def bottleNeck1(input, inchannel, outchannel):
    x = tf.nn.conv2d(input, filter=[1, 1, inchannel, inchannel], strides=[1, 1, 1, 1], padding='valid')

    x = tf.nn.conv2d(x, strides=[1, 1, 1, 1])



def correctPad(img_input, kernelsize):  #:算法构建填充尺寸
    if isinstance(kernelsize, int):
        kernelsize = (kernelsize, kernelsize)
    inputShape = tf.keras.backend.int_shape(img_input)[1:3]
    if inputShape[0] is None:
        adjust = (1, 1)
    else:  # 如果输入尺寸为偶数则需加0，
        adjust = (1 - inputShape[0] % 2, 1 - inputShape[1] % 2)
    corr = (kernelsize[0] // 2, kernelsize[1] // 2)
    return ((corr[0] - adjust[0], corr[0]),
            (corr[1] - corr[1], corr[1]))


def BottleNeck(input, expansion=6, out_channel=16, stride=1, padding=None):
    inChannel = tf.keras.backend.int_shape(input)[-1]
    # TODO:outChannel通过alpha参数计算得出，不指定
    outChannel = out_channel
    expanChannel = expansion * inChannel
    # bottleneck第一步
    x = input
    x = tf.keras.layers.Conv2D(expanChannel, kernel_size=(1, 1),
                               padding='same',
                               use_bias=False)(x)  # 112*112*32
    x = tf.keras.layers.BatchNormalization(axis=3)(x)
    x = tf.keras.layers.ReLU(max_value=6.)(x)

    if stride == 2:
        x = tf.keras.layers.ZeroPadding2D(padding=correctPad(x, 3))
    x = tf.keras.layers.DepthwiseConv2D(kernel_size=(3, 3),
                                        strides=stride,
                                        padding='same' if stride == 1 else "valid")(x)
    x = tf.keras.layers.BatchNormalization(axis=3)(x)
    x = tf.keras.layers.ReLU(max_value=6.)(x)

    # Project
    x = tf.keras.layers.Conv2D(outChannel,
                               kernel_size=(1, 1),
                               padding='same',
                               use_bias=False)(x)
    x = tf.keras.layers.BatchNormalization()(x)
    if stride == 1:
        return tf.keras.layers.Add()([input, x])
    return x


# 编写一个或多个数据集导入函数 #another:tf.estimator.inputs.pandas_input_fn(train_df, train_df["polarity"], num_epochs=None, shuffle=True)
def trainInputFn(features, labels, batch_size, epoch):
    dataset = tf.data.Dataset.from_tensor_slices((dict(features), labels))
    ds = dataset.shuffle(1000).repeat(epoch).batch(batch_size)
    # 生成批次(features, labels)对的输入管道
    return


# 定义特征列
# Feature columns describe how to use the input.
img_feature_column = tf.feature_column.numeric_column(
    key='img',
    dtype=tf.uint8
)


def myModelFn(features, labels, mode, params):
    """
    :param features: This is batch_features from input_fn
    :param labels:This is batch_features from input_fn
    :param mode:An instance of tf.estimator.ModeKeys
    :param params:Additional configuration
    :return:
    """
    # 针对图像
    print('features', features.head())
    exit(0)
    base_input = tf.reshape(features, (-1, 224, 224, 3))
    # FIXME:重构代码

    x = tf.keras.layers.ZeroPadding2D(padding=correctPad(img_input=base_input, kernelsize=3))(base_input)  # 225*225*3
    x = tf.keras.layers.Conv2D(32, (3, 3), strides=(2, 2), use_bias=False)(x)  # 112*112*32
    x = tf.keras.layers.BatchNormalization(axis=3)(x)
    x = tf.keras.layers.ReLU(max_value=6)(x)

    x = BottleNeck(x, expansion=1, out_channel=16, stride=1)  # 112*112*16

    x = BottleNeck(x, expansion=6, out_channel=24, stride=2)  # 56*56*24
    x = BottleNeck(x, expansion=6, out_channel=24, stride=1)  # 56*56*24

    x = BottleNeck(x, expansion=6, out_channel=32, stride=2)  # 28*28*32
    x = BottleNeck(x, expansion=6, out_channel=32, stride=1)
    x = BottleNeck(x, expansion=6, out_channel=32, stride=1)

    x = BottleNeck(x, expansion=6, out_channel=64, stride=2)  # 14*14*64
    x = BottleNeck(x, expansion=6, out_channel=64, stride=1)
    x = BottleNeck(x, expansion=6, out_channel=64, stride=1)
    x = BottleNeck(x, expansion=6, out_channel=64, stride=1)

    x = BottleNeck(x, expansion=6, out_channel=96, stride=1)  # 14*14*96
    x = BottleNeck(x, expansion=6, out_channel=96, stride=1)
    x = BottleNeck(x, expansion=6, out_channel=96, stride=1)

    x = BottleNeck(x, expansion=6, out_channel=160, stride=2)  # 7*7*160
    x = BottleNeck(x, expansion=6, out_channel=160, stride=1)

    x = BottleNeck(x, expansion=6, out_channel=320, stride=1)  # 7*7*320

    x = tf.keras.layers.Conv2D(1280, kernel_size=(1, 1), strides=(1, 1))(x)
    x = tf.keras.layers.GlobalAveragePooling2D()(x)
    logits = tf.keras.layers.Dense(params.get('classes'))(x)

    # 预测
    predictions = {
        'classes': tf.argmax(logits, axis=1),
        'probabilities': tf.keras.layers.Softmax()(logits)
    }
    if mode == tf.estimator.ModeKeys.PREDICT:
        return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)

    loss = tf.losses.sparse_softmax_cross_entropy(labels, logits)
    if mode == tf.estimator.ModeKeys.TRAIN:
        optimizer = tf.train.AdagradOptimizer(learning_rate=1e-3)
        train_op = optimizer.minimize(loss,
                                      global_step=tf.train.get_global_step())
        return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)
    # Add evaluation metrics (for EVAL mode)
    eval_metric_ops = {
        'accuracy': tf.metrics.accuracy(labels=labels,
                                        predictions=predictions["classes"])
    }
    return tf.estimator.EstimatorSpec(mode=mode, loss=loss, eval_metric_ops=eval_metric_ops)


if __name__ == "__main__":
    totals = 150528
    columns = []
    for i in range(totals):
        columns.append(f'pixel{i + 1}')
    columns.append('label')
    ckptPath = r'F:\Resources\backup\fatp'
    myModel = tf.estimator.Estimator(
        model_fn=myModelFn,
        model_dir=ckptPath,
        params={
            'classes': 8
        }
    )
    csvFile = r"F:\Resources\backup\img.csv"
    # TODO:csv 行解析器
    trainDf = pd.read_csv(csvFile, header=None, names=columns, skiprows=1)
    label = trainDf.pop("label")
    print(trainDf.head())
    print('read data done')
    # 这是一种定义输入函数就把数据也传输了

    # trainInputFn = functools.partial(trainInputFn,features={"img":trainDf},
    #                                  labels=label,
    #                                  batch_size=32,
    #                                  epoch=200)
    train_input_fn = tf.estimator.inputs.pandas_input_fn(
        trainDf,
        label,
        batch_size=100,
        num_epochs=2,
        shuffle=True)
    print('train input funcion done')
    myModel.train(
        input_fn=train_input_fn,
    )
